"""
回测引擎
支持对历史K线数据进行策略回测
"""
import pandas as pd
import numpy as np
from typing import Dict, List, Optional
from datetime import datetime
from src.data.processor import KlineProcessor
from src.strategy.scalping import ScalpingStrategy
from src.trading.position import PositionManager
from src.trading.risk import RiskManager, RiskControl
from src.trading.exit_strategy import ExitStrategy
from src.costs.calculator import CostCalculator
from src.utils.logger import get_logger

logger = get_logger(__name__)


class BacktestEngine:
    """回测引擎"""

    def __init__(self, config: Dict, initial_balance: float = 10000):
        """
        初始化回测引擎

        Args:
            config: 策略配置
            initial_balance: 初始资金
        """
        self.config = config
        self.initial_balance = initial_balance
        self.current_balance = initial_balance

        # 初始化组件
        self.processor = KlineProcessor()
        self.strategy = ScalpingStrategy(config.get('strategy', {}))
        self.position_manager = PositionManager()
        self.exit_strategy = ExitStrategy(config.get('exit', {}))
        self.cost_calculator = CostCalculator(
            maker_fee=config.get('costs', {}).get('maker_fee', 0.02),
            taker_fee=config.get('costs', {}).get('taker_fee', 0.05),
            funding_rate=config.get('costs', {}).get('funding_rate', 0.03)
        )

        # 初始化风险管理
        risk_control = RiskControl(
            stop_loss_pct=config.get('risk', {}).get('stop_loss', 1.0),
            take_profit_pct=config.get('risk', {}).get('take_profit', 2.0),
            max_positions=config.get('risk', {}).get('max_positions', 3),
            max_daily_loss_pct=config.get('risk', {}).get('max_daily_loss', 5.0),
            max_drawdown_pct=config.get('risk', {}).get('max_drawdown', 10.0),
            max_position_bars=config.get('trading', {}).get('max_position_bars', 20),
        )
        self.risk_manager = RiskManager(risk_control, initial_balance)

        # 回测统计
        self.stats = {
            'bars_processed': 0,
            'trades_executed': 0,
            'peak_balance': initial_balance,
            'lowest_balance': initial_balance,
        }

    def run_backtest(self, symbols: List[str], klines_data: Dict[str, Dict[int, List]],
                     start_date: datetime = None, end_date: datetime = None) -> Dict:
        """
        执行回测

        Args:
            symbols: 交易对列表
            klines_data: {symbol: {timeframe: [klines]}} 格式的K线数据
            start_date: 开始日期
            end_date: 结束日期

        Returns:
            回测结果字典
        """
        logger.info(f"开始回测: 交易对={symbols}, 初始资金={self.initial_balance}")

        # 处理K线数据
        for symbol in symbols:
            if symbol not in klines_data:
                logger.warning(f"跳过缺失数据的交易对: {symbol}")
                continue

            for timeframe, klines in klines_data[symbol].items():
                # 转换为DataFrame
                data = {
                    'timestamp': [k.timestamp for k in klines],
                    'open': [k.open for k in klines],
                    'high': [k.high for k in klines],
                    'low': [k.low for k in klines],
                    'close': [k.close for k in klines],
                    'volume': [k.volume for k in klines],
                }
                df = pd.DataFrame(data)

                # 添加技术指标
                df = self.processor.add_technical_indicators(df, self.config.get('strategy', {}))

                # 存储处理后的数据
                self.processor.add_klines(symbol, timeframe, klines)

        # 回测主循环
        for symbol in symbols:
            df_5m = self.processor.get_klines(symbol, 5)
            if df_5m is None or len(df_5m) < 20:
                logger.warning(f"数据不足，跳过: {symbol}")
                continue

            self._backtest_symbol(symbol, df_5m)

        # 生成回测报告
        report = self._generate_report()
        return report

    def _backtest_symbol(self, symbol: str, df_5m: pd.DataFrame):
        """
        对单个交易对进行回测

        Args:
            symbol: 交易对
            df_5m: 5分钟K线DataFrame
        """
        logger.info(f"开始回测: {symbol}, 共{len(df_5m)}条K线")

        # 获取其他时间框架的数据
        df_15m = self.processor.get_klines(symbol, 15)
        df_1h = self.processor.get_klines(symbol, 60)
        df_4h = self.processor.get_klines(symbol, 240)

        # 记录每个持仓的开仓K线索引 (用于计算持仓时长)
        position_entry_bar = {}

        # 逐K线回测
        for i in range(20, len(df_5m)):
            current_bar = df_5m.iloc[i]
            current_time = pd.to_datetime(current_bar['timestamp'], unit='ms')

            # 收集各时间框架最新的数据
            kline_data = {
                5: df_5m.iloc[i-19:i+1].copy(),
            }

            if df_15m is not None and len(df_15m) > 0:
                kline_data[15] = df_15m.tail(20).copy()

            if df_1h is not None and len(df_1h) > 0:
                kline_data[60] = df_1h.tail(20).copy()

            if df_4h is not None and len(df_4h) > 0:
                kline_data[240] = df_4h.tail(20).copy()

            # 检查现有持仓
            position = self.position_manager.get_active_position(symbol)

            # 平仓逻辑
            if position:
                # 计算持仓K线数
                bars_held = i - position_entry_bar.get(position.trade_id, i)

                # 使用新的出场策略
                should_exit, reason = self.exit_strategy.should_exit(
                    position, current_bar, kline_data, bars_held,
                    params=self.config.get('risk', {})
                )

                if should_exit:
                    self._close_position(position, current_bar, reason)
                    if position.trade_id in position_entry_bar:
                        del position_entry_bar[position.trade_id]
            else:
                # 生成交易信号
                signal, confidence = self.strategy.generate_signal(kline_data)

                if signal in ["buy", "sell"] and confidence > 0.6:
                    trade = self._open_position(symbol, signal, current_bar, confidence)
                    if trade:
                        position_entry_bar[trade.trade_id] = i

            self.stats['bars_processed'] += 1

    def _open_position(self, symbol: str, signal: str, current_bar: pd.Series,
                      confidence: float):
        """开仓"""
        # 检查风险
        position_size = self.current_balance * 0.05  # 每笔5%的资金

        if not self.risk_manager.can_open_position(position_size):
            return None

        price = current_bar['close']

        trade = self.position_manager.open_position(
            symbol=symbol,
            side="long" if signal == "buy" else "short",
            price=price,
            quantity=position_size / price,
            timestamp=pd.to_datetime(current_bar['timestamp'], unit='ms').to_pydatetime()
        )

        self.stats['trades_executed'] += 1
        logger.debug(f"[开仓] {symbol} {signal.upper()}, 价格={price:.2f}, 信心度={confidence:.2f}")

        return trade

    def _close_position(self, position, current_bar: pd.Series, reason: str):
        """平仓"""
        price = current_bar['close']

        # 计算盈亏
        if position.side == "long":
            pnl = (price - position.entry_price) * position.entry_qty
        else:
            pnl = (position.entry_price - price) * position.entry_qty

        # 扣除费用
        cost_info = self.cost_calculator.calculate_profit_loss(
            entry_price=position.entry_price,
            exit_price=price,
            quantity=position.entry_qty,
            trade_type=self.config.get('trading', {}).get('trade_type', 'spot')
        )

        pnl_after_cost = cost_info['net_pnl']

        # 平仓
        closed_trade = self.position_manager.close_position(
            trade_id=position.trade_id,
            exit_price=price,
            exit_qty=position.entry_qty,
            timestamp=pd.to_datetime(current_bar['timestamp'], unit='ms').to_pydatetime(),
            reason=reason,
            cost_info=cost_info
        )

        # 更新账户余额
        self.current_balance += pnl_after_cost
        self.risk_manager.update_balance(self.current_balance)

        # 更新统计
        self.stats['peak_balance'] = max(self.stats['peak_balance'], self.current_balance)
        self.stats['lowest_balance'] = min(self.stats['lowest_balance'], self.current_balance)

        logger.debug(f"[平仓] {position.symbol} {reason}, 盈亏={pnl_after_cost:.2f}, "
                    f"账户余额={self.current_balance:.2f}")

    def _calculate_pnl_pct(self, side: str, entry_price: float, current_price: float) -> float:
        """计算盈亏百分比"""
        if side == "long":
            return ((current_price - entry_price) / entry_price) * 100
        else:
            return ((entry_price - current_price) / entry_price) * 100

    def _generate_report(self) -> Dict:
        """生成回测报告"""
        trades = self.position_manager.get_closed_trades()
        stats = self.position_manager.calculate_stats()

        total_return = self.current_balance - self.initial_balance
        return_pct = (total_return / self.initial_balance) * 100

        max_drawdown = (
            (self.stats['lowest_balance'] - self.stats['peak_balance']) / self.stats['peak_balance']
        ) * 100

        report = {
            'initial_balance': self.initial_balance,
            'final_balance': self.current_balance,
            'total_return': total_return,
            'return_pct': return_pct,
            'max_drawdown_pct': max_drawdown,
            'bars_processed': self.stats['bars_processed'],
            'trades': {
                'total': stats['total_trades'],
                'winning': stats['winning_trades'],
                'losing': stats['losing_trades'],
                'win_rate_pct': stats['win_rate'],
            },
            'pnl': {
                'total': stats['total_pnl'],
                'avg_per_trade': stats['avg_pnl'],
                'avg_win': stats['avg_win'],
                'avg_loss': stats['avg_loss'],
                'profit_factor': stats['profit_factor'],
            }
        }

        logger.info(f"\n{'='*50}")
        logger.info("回测报告")
        logger.info(f"{'='*50}")
        logger.info(f"初始资金: ${self.initial_balance:.2f}")
        logger.info(f"最终资金: ${self.current_balance:.2f}")
        logger.info(f"总收益: ${total_return:.2f} ({return_pct:.2f}%)")
        logger.info(f"最大回撤: {max_drawdown:.2f}%")
        logger.info(f"处理K线: {self.stats['bars_processed']}")
        logger.info(f"总交易数: {stats['total_trades']}")
        logger.info(f"胜率: {stats['win_rate']:.2f}%")
        logger.info(f"平均单笔盈亏: ${stats['avg_pnl']:.2f}")
        logger.info(f"盈利因子: {stats['profit_factor']:.2f}")
        logger.info(f"{'='*50}\n")

        return report
